#!/usr/bin/env python3
import os
import subprocess
from pathlib import Path
from datetime import datetime, timedelta

class CertificateGenerator:
    def __init__(self, cert_dir="./certs"):
        self.cert_dir = Path(cert_dir)
        self.cert_dir.mkdir(exist_ok=True)
        self.ca_key = self.cert_dir / "ca.key"
        self.ca_cert = self.cert_dir / "ca.crt"
        
    def create_ca_certificate(self):
        """Create Certificate Authority for HTTPS interception"""
        if self.ca_key.exists() and self.ca_cert.exists():
            print("CA certificate already exists")
            return True
            
        try:
            # Generate CA private key
            cmd_key = [
                "openssl", "genrsa", 
                "-out", str(self.ca_key), 
                "2048"
            ]
            subprocess.run(cmd_key, check=True, capture_output=True)
            
            # Generate CA certificate
            cmd_cert = [
                "openssl", "req", "-new", "-x509",
                "-key", str(self.ca_key),
                "-out", str(self.ca_cert),
                "-days", "365",
                "-subj", "/C=US/ST=Test/L=Test/O=IDEA Proxy/OU=Security/CN=IDEA Proxy CA"
            ]
            subprocess.run(cmd_cert, check=True, capture_output=True)
            
            print(f"CA certificate created: {self.ca_cert}")
            print(f"CA private key created: {self.ca_key}")
            return True
            
        except subprocess.CalledProcessError as e:
            print(f"Error creating CA certificate: {e}")
            return False
    
    def create_server_certificate(self, hostname):
        """Create server certificate for specific hostname"""
        cert_file = self.cert_dir / f"{hostname}.crt"
        key_file = self.cert_dir / f"{hostname}.key"
        
        if cert_file.exists() and key_file.exists():
            return str(cert_file), str(key_file)
        
        try:
            # Generate server private key
            cmd_key = [
                "openssl", "genrsa",
                "-out", str(key_file),
                "2048"
            ]
            subprocess.run(cmd_key, check=True, capture_output=True)
            
            # Generate certificate signing request
            csr_file = self.cert_dir / f"{hostname}.csr"
            cmd_csr = [
                "openssl", "req", "-new",
                "-key", str(key_file),
                "-out", str(csr_file),
                "-subj", f"/C=US/ST=Test/L=Test/O=IDEA Proxy/OU=Security/CN={hostname}"
            ]
            subprocess.run(cmd_csr, check=True, capture_output=True)
            
            # Generate server certificate signed by CA
            cmd_cert = [
                "openssl", "x509", "-req",
                "-in", str(csr_file),
                "-CA", str(self.ca_cert),
                "-CAkey", str(self.ca_key),
                "-CAcreateserial",
                "-out", str(cert_file),
                "-days", "30"
            ]
            subprocess.run(cmd_cert, check=True, capture_output=True)
            
            # Clean up CSR file
            csr_file.unlink()
            
            print(f"Certificate created for {hostname}: {cert_file}")
            return str(cert_file), str(key_file)
            
        except subprocess.CalledProcessError as e:
            print(f"Error creating certificate for {hostname}: {e}")
            return None, None
    
    def get_install_instructions(self):
        """Get instructions for installing CA certificate"""
        instructions = f"""
To install the CA certificate for HTTPS interception:

1. Linux (system-wide):
   sudo cp {self.ca_cert} /usr/local/share/ca-certificates/idea-proxy-ca.crt
   sudo update-ca-certificates

2. macOS:
   sudo security add-trusted-cert -d -r trustRoot -k /Library/Keychains/System.keychain {self.ca_cert}

3. Windows:
   certmgr.msc -> Trusted Root Certification Authorities -> Certificates -> Import {self.ca_cert}

4. Firefox (manual):
   Settings -> Privacy & Security -> Certificates -> View Certificates -> Import {self.ca_cert}

5. IDEA (if needed):
   File -> Settings -> Tools -> Server Certificates -> Import {self.ca_cert}
"""
        return instructions

if __name__ == '__main__':
    generator = CertificateGenerator()
    
    if generator.create_ca_certificate():
        print("\nCA Certificate created successfully!")
        print(generator.get_install_instructions())
        
        # Test creating a certificate for common hosts
        test_hosts = ['jetbrains.com', 'plugins.jetbrains.com', 'www.jetbrains.com']
        for host in test_hosts:
            cert, key = generator.create_server_certificate(host)
            if cert and key:
                print(f"Test certificate created for {host}")
    else:
        print("Failed to create CA certificate")